import kline
import pandas as pd
import talib
import numpy as np
from colorama import init, Fore, Back, Style
import pivot as pvt
import mplfinance as mpf
import matplotlib.pyplot as plt

K_TYPE_UP = "A"  # 上涨
K_TYPE_DOWN = "B"  # 下跌
K_TYPE_EXPAND_N = "C"  # 正常板块扩张
K_TYPE_EXPAND_UP = "D"  # 上扩张
K_TYPE_EXPAND_DOWN = "E"  # 下扩张
K_TYPE_RETRACT_N = "F"  # 正常收敛
K_TYPE_RETRACT_UP = "G"  # 上收敛
K_TYPE_RETRACT_DOWN = "H"  # 下收敛
K_TYPE_SMOOTH_N = "I"  # 正常平稳
K_TYPE_JUMP_UP = "J"  # 上跳空
K_TYPE_JUMP_DOWN = "K"  # 下跳空
K_TYPE_UNKNOWN = "Z"  # 未知
MACD_START = 50  # MACD计算开始位置
MACD_THRESHOLD = 0.05  # MACD的阈值

# 根据两个K线判断规则,type=类型，direction=方向 1:向上 2:向下 3:平
def judge2k(k1, k2):
    res={}
    # 判断上涨
    if k1.low < k2.low and k1.high < k2.high and k2.low <= k1.high:
        res['type']=K_TYPE_UP
        res['direction']=1
    # 判断下跌
    elif k1.high > k2.high and k1.low > k2.low and k2.high >= k1.low:
        res['type']=K_TYPE_DOWN
        res['direction']=2
    # 判断正常板块扩张
    elif k1.high < k2.high and k2.low < k1.low:
        res['type']=K_TYPE_EXPAND_N
        res['direction']=3
    # 判断上扩张
    elif k1.high < k2.high and k2.low == k1.low:
        res['type']=K_TYPE_EXPAND_UP
        res['direction']=3
    # 判断下扩张
    elif k1.low > k2.low and k2.high == k1.high:
        res['type']=K_TYPE_EXPAND_DOWN
        res['direction']=3
    # 判断正常收敛
    elif k1.high > k2.high and k2.low > k1.low:
        res['type']=K_TYPE_RETRACT_N
        res['direction']=3
    # 判断上收敛
    elif k1.high == k2.high and k2.low > k1.low:
        res['type']=K_TYPE_RETRACT_UP
        res['direction']=3
    # 判断下收敛
    elif k1.low == k2.low and k2.high < k1.high:
        res['type']=K_TYPE_RETRACT_DOWN
        res['direction']=3
    # 判断正常平稳
    elif k1.low == k2.low and k2.high == k1.high:
        res['type']=K_TYPE_SMOOTH_N
        res['direction']=3
    # 判断上跳空
    elif k2.low > k1.high:
        res['type']=K_TYPE_JUMP_UP
        res['direction']=1
    # 判断下跳空
    elif k2.high < k1.low:
        res['type']=K_TYPE_JUMP_DOWN
        res['direction']=2
    # 未知
    else:
        res['type']=K_TYPE_UNKNOWN
        res['direction']=3
        
    # print('对比结果：',res)
    return res
    

#根据K线列表，返回macd信息,目前用DIF值，所以只返回MACD值
def get_macd(k_list):
    if len(k_list) < MACD_START:
        return []
    #计算MACD
    else:
        #macd=dif,signal=dea,hist=macd柱状图
        macd, signal, hist = talib.MACD(k_list['close'], fastperiod=12, slowperiod=26, signalperiod=9)
        return macd
    
#吧所有K点，转换成折线信息,因为要取MACD值，所以K线数量要大于50，且从第50条开始计算，因为MACD前几十条没有值
#连续上涨，连续下跌，连续平都算一根。
#连续平，则判断是否有中枢
#入参，k_list=K线列表，line3_type=1表示平线中枢按照最大区间判断，=2时表示按公共区间判断
def k_to_line(k_list={},line3_type=1):
    res = []
    last_k ={}#如果当前为平线时，需要重新处理平线，重新判断高低值。  本值用来保存最后一根K线的信息
    line_start=-1#折线开始位置
    line_end=0#折线结束位置
    line_direction=0#当前走势的方向
    line_max=0#折线的最高点
    line_min=0#折线的最低点
    line_num=0#折线中包含K线数量
    line_jump_time=0#跳空次数
    line_ifzs=0#是否形成中枢
    pivot_high=0#中枢区间的最高点
    pivot_low=0#中枢区间的最低点
    line = kline.Kline(line_start, line_end, line_direction, line_max, line_min, line_num, line_jump_time,pivot_high=pivot_high,pivot_low=pivot_low)
    #目前只用dif，且数据不用处理，所以直接调talib方法
    macd, signal, hist = talib.MACD(k_list['close'], fastperiod=12, slowperiod=26, signalperiod=9)
    macd=np.array(macd)#重新排列数组下标，否则不是从0开始
    if len(macd) < 2:
        return res
    if len(k_list) < 2:
        return res
    else:
        for i in range(MACD_START,len(k_list) - 1):
            k1 = k_list.iloc[i]
            k2 = k_list.iloc[i + 1]
            #判断当前是否为平线，如果为平线，则将last_k的值赋值给k1
            if line_direction == 3 and line3_type == 2:
                print('平线判断：',last_k)
                k1.high = last_k['high']
                k1.low = last_k['low']
            rule = judge2k(k1, k2)    
            # print('K线对比：',i,k1,k2,rule)
            line_num=line_num+1
            #没有方向=初始情况，先赋值
            if line_direction == 0 :
                line_start = i
                line_stime = k1['day']
                line_direction = rule['direction']
                if line_direction == 3 and line3_type == 2:
                    line_max = min(line_max, float(k1.high), float(k2.high))
                    line_min = max(line_min, float(k1.low), float(k2.low))
                    last_k['high'] = line_max
                    last_k['low'] = line_min
                else:
                    line_max = float(max(k1.high, k2.high))
                    line_min = float(min(k1.low, k2.low))
            #相同方向,判断最高，最低点(转折不能在这判断)
            if line_direction == rule['direction'] :
                if line_direction == 3 and line3_type == 2:
                    line_max = min(line_max, float(k1.high), float(k2.high))
                    line_min = max(line_min, float(k1.low), float(k2.low))
                    last_k['high'] = line_max
                    last_k['low'] = line_min
                else:
                    line_max = max(line_max, float(k2.high))
                    line_min = min(line_min, float(k2.low))        
            
            #方向不同，产生折叠
            if line_direction != rule['direction'] and line_direction != 0:
                
                #将之前的折线记录到结果中
                line_end = i
                line_num=line_num-1
                line.line_start = line_start
                line.line_end = line_end
                line.line_direction = line_direction
                line.line_max = line_max
                line.line_min = line_min
                line.line_jump_time = line_jump_time
                line.line_num = line_num
                line.line_stime = line_stime
                line.line_endtime = k1['day']
                line.line_ifzs = line_ifzs
                line.pivot_high = pivot_high
                line.pivot_low = pivot_low
                res.append(line)
                
                # print('出现转折：',i,vars(line))
                
                #将新折线信息初始化
                line_start = i
                line_stime = k1['day']
                line_direction = rule['direction']
                line_jump_time = 0
                line_num = 1
                if line_direction == 3 and line3_type == 2:
                    line_max = min(line_max, float(k1.high), float(k2.high))
                    line_min = max(line_min, float(k1.low), float(k2.low))
                    last_k['high'] = line_max
                    last_k['low'] = line_min
                else:
                    line_max = float(max(k1.high, k2.high))
                    line_min = float(min(k1.low, k2.low))
                line_ifzs = 0
                pivot_high = 0
                pivot_low = 0
                line = kline.Kline(line_start, line_end, line_direction, line_max, line_min, line_num, line_jump_time,line_ifzs=line_ifzs,pivot_high=pivot_high,pivot_low=pivot_low)
            #这里需要放在z轴判断之后，因为要考虑属于新的还是老的折线
            #如果是平，之前没有中枢，则判断是否产生中枢
            if  rule['direction'] == 3:
                if line_ifzs == 0:       
                #判断是否形成中枢
                    if macd[i]<MACD_THRESHOLD:
                        line_ifzs=1
                        if line3_type == 1:
                            pivot_high = max(float(k1.high), float(k2.high))
                            pivot_low = min(float(k1.low), float(k2.low))
                        else:
                            pivot_high = min(float(k1.high), float(k2.high))
                            pivot_low = max(float(k1.low), float(k2.low))
                        print('平线可形成中枢：',i)            
           
            #判断跳空
            if rule['type'] in [K_TYPE_JUMP_UP, K_TYPE_JUMP_DOWN]:
                line_jump_time += 1          
                
            #判断最后一根K线
            if i == len(k_list) - 2:
                line_end = i + 1
                line.line_start = line_start
                line.line_end = line_end
                line.line_direction = line_direction
                line.line_max = line_max
                line.line_min = line_min
                line.line_jump_time = line_jump_time
                line.line_num = line_num
                line.line_stime = k1['day']
                line.line_endtime = k2['day']
                res.append(line)
    
    return res

#画折线图
def draw_foldline(line_list):
    for line in line_list:
        if line.line_direction == 1:
            plt.plot([line.line_stime,line.line_endtime],[line.line_min,line.line_max],color='red')
        if line.line_direction == 2:
            plt.plot([line.line_stime,line.line_endtime],[line.line_max,line.line_min],color='green')
        if line.line_direction == 3:
            plt.plot([line.line_stime,line.line_endtime],[line.line_min,line.line_min],color='gray')
            plt.plot([line.line_stime,line.line_endtime],[line.line_max,line.line_max],color='gray')
            
#画中枢图
def draw_pivot(pivot_list):
    
    for pivot in pivot_list:
        plt.plot([pivot.pivot_stime,pivot.pivot_endtime],[pivot.pivot_low,pivot.pivot_low],color='#CC3366',linestyle='--',label='pdown')
        plt.plot([pivot.pivot_stime,pivot.pivot_endtime],[pivot.pivot_high,pivot.pivot_high],color='#CC3366',linestyle='-.',label='pup')



#判断这条线是否结束中枢，此时line一定不是平线，且一定是转折
#0=没结束，1=结束
def if_zs_break(line,pivot,last_line):
    res = 0
    #如果最后一根线和本次的线都不在中枢区间，则说明中枢结束
    if judge_line_inrange(last_line,pivot) == 0 and judge_line_inrange(line,pivot) == 0:
        res = 1
    
    return res

#判断这条线是否在中枢区间
#1:在区间，0：不在区间
def judge_line_inrange(line,pivot):
    res = 1 
    if line.line_max < pivot.pivot_low or line.line_min > pivot.pivot_high:
        res = 0
    return res

#判断这几条线的数组是否形成中枢
#返回：是否形成中枢，中枢区间最低点，中枢区间最高点
def if_zs_create(trend):
    res = 0
    pivot_low = 0
    pivot_high = 0
    start=0#起始方向
    line_s = {}
    line_e = {}
    #找到第一根和最后一根线，判断区间
    for line in trend:
        if line.line_direction == 3:
            continue
        if start == 0:
            start = line.line_direction
            line_s = line
        elif line.line_direction == start:
            line_e = line
            break
    if line_s == {} or line_e == {}:
        return res,0,0
    #判断方向
    #上升方式
    if start == 1:
        if line_s.line_min < line_e.line_max:
            res = 1
            pivot_low = max(line_s.line_min, line_e.line_min)
            pivot_high = min(line_s.line_max, line_e.line_max)
    elif start == 2:
        if line_s.line_max > line_e.line_min:
            res = 1
            pivot_low = max(line_s.line_min, line_e.line_min)
            pivot_high = min(line_s.line_max, line_e.line_max)
    
    return res,pivot_low,pivot_high

#根据折线列表，生成中枢情况
def line_to_pivot(line_list):
    res = []
    if len(line_list) < 3:
        return res
    pivot_start = -1#z轴开始位置
    pivot_end = 0#z轴结束位置
    pivot_stime = 0#中枢开始时间
    pivot_endtime = 0#中枢结束时间
    pivot_max = 0#中枢区间的最高点
    pivot_min = 0#中枢区间的最低点
    pivot_high = 0#中枢震荡区间的最高点
    pivot_low = 0#中枢震荡区间的最低点
    pivot_type = -1#中枢类型，1=折线中枢，2=平线中枢
    pivot_jump_time = 0#中枢跳空次数
    pivot_trend = []#中枢走势
    pivot = {}
    last_line = {}#记录最后一根线的信息（同方向的整根线）
    ifzs=0#记录当前是否为有中枢状态
    
    for i in range(len(line_list)):
        
        #判断当前是否已经形成中枢
        if ifzs == 1:
            
            line = line_list[i]
            #计算最后一根线,只考虑非平情况
            if line.line_direction != 3:
                if last_line == {}:
                    last_line = line
                if line.line_direction == last_line.line_direction:
                    last_line.line_max = max(last_line.line_max, line.line_max)
                    last_line.line_min = min(last_line.line_min, line.line_min)
                else:
                    #判断本次曲线是否导致中枢结束
                    # print('判断中枢：',i,vars(pivot))
                    # print('判断中枢——线：',vars(line))
                    # print('判断中枢-最后一根线：',vars(last_line))
                    if if_zs_break(line,pivot,last_line) == 1:
                        ifzs = 0
                        pivot_end = i
                        pivot_endtime = line_list[i].line_endtime
                        pivot.pivot_end = pivot_end
                        pivot.pivot_endtime = pivot_endtime
                        res.append(pivot)
                        pivot = {}
                        print('中枢结束：',i,pivot)
                    else:
                        #重新赋值已形成的中枢，计算最大值，最小值，折线次数等
                        pivot.pivot_max = max(pivot.pivot_max, line.line_max)
                        pivot.pivot_min = min(pivot.pivot_min, line.line_min)
                        pivot.pivot_trend.append(line)
                    #把本线段变成最后一根线
                    last_line = line
            
            
                
            
        #如果没形成中枢.往后查3跟折线计算
        else:
            linei = line_list[i]
            trend = []#记录折线走势
            trend3 = []#记录3次折线走势，用来判断是否形成中枢时用
            direction_line ={}#记录当前方向的线
            fold_time = 0#折叠次数，只有上下改变才算折叠，平线不算
            direction = 0#记录当前走势的方向
            j=i
            while fold_time<=3:
                if j>=len(line_list):
                    break
                line = line_list[j]
                
                pivot_max = max(pivot_max, line.line_max)
                pivot_min = min(pivot_min, line.line_min)
                #判断是否为平线
                if line.line_direction == 3:
                    #判断是否形成中枢
                    if line.line_ifzs==1:
                        ifzs = 1
                        i=j
                        pivot_start = j
                        pivot_stime = line.line_stime
                        pivot_type = 2
                        pivot_high = line.pivot_high
                        pivot_low = line.pivot_low
                        pivot_max = line.line_max
                        pivot_min = line.line_min
                        pivot_trend = [line]
                        pivot = pvt.pivot(pivot_start, pivot_end, pivot_stime, pivot_endtime, pivot_max, pivot_min, pivot_high, pivot_low,pivot_type, pivot_trend, fold_time)
                        fold_time = 0
                        print('形成平线中枢：',i,vars(pivot))
                        break
                    else:
                        trend.append(line)      
                           
                #正常折线
                else:       
                    #判断为异向
                    if direction != line.line_direction:
                        if direction_line!= {}:
                            trend3.append(direction_line)
                        direction = line.line_direction
                        fold_time = fold_time+1
                        direction_line = line
                    #判断为同向
                    else:
                        direction_line.line_max = max(direction_line.line_max, line.line_max)
                        direction_line.line_min = min(direction_line.line_min, line.line_min)
                    trend.append(line)      
                    j=j+1
                    #当fold_time=4时，表示凑够3跟折线，且折线已经结束，开始判断是否形成中枢
                    if fold_time == 4:
                        zsres,pivot_low,pivot_high = if_zs_create(trend3)
                        if zsres == 1:
                            #形成中枢，开始赋值
                            ifzs = 1
                            pivot_start = i
                            pivot_stime = linei.line_stime
                            pivot_type = 1
                            pivot_high = pivot_high
                            pivot_low = pivot_low
                            pivot_max = pivot_max
                            pivot_min = pivot_min
                            pivot_trend = trend
                            pivot = pvt.pivot(pivot_start, pivot_end, pivot_stime, pivot_endtime, pivot_max, pivot_min, pivot_high, pivot_low,pivot_type, pivot_trend, fold_time)
                            print('形成中枢：',i,vars(pivot))
                            i=j
                            break
                    

    if i==len(line_list)-1 and ifzs == 1:
        pivot_end = i
        pivot_endtime = line_list[i].line_endtime
        pivot.pivot_end = pivot_end
        pivot.pivot_endtime = pivot_endtime
        res.append(pivot)
        print('中枢结束：',i,vars(pivot))
    return res
                
                
        
            

# #根据列表判断规则
def judge_k_list(k_list={},line3_type=1):
    res={}
    
    #先将日期转化为日期格式
    k_list['date'] = pd.to_datetime(k_list['day'])
    k_list.set_index('date',inplace=True)
    #将K线变成折线
    line_list = k_to_line(k_list,line3_type)
    
    #画折线图
    plt.figure()
    draw_foldline(line_list)
    # print("共有折线:",len(line_list))
    # print("折线列表:")
    # for line in line_list:
    #     if line.line_direction == 1:
    #         print(Fore.RED + str(vars(line)))
    #     elif line.line_direction == 2:
    #         print(Fore.GREEN + str(vars(line)))
    #     else:
    #         print(Fore.WHITE +str(vars(line)))
    pivot_list = line_to_pivot(line_list)
    print("中枢列表:")
    for pivot in pivot_list:
        print(Fore.YELLOW + str(vars(pivot)))
    
    #画中枢图
    draw_pivot(pivot_list)
    
    #画K线图
    mpf.plot(k_list[MACD_START:],type='candle',style='charles',title='K线图')
    plt.show()
    
    # if len(line_list) < 3:
    #     return res
